from enhanced_classifier import EnhancedClassifier
from visualization import ClassificationVisualizer
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import label_binarize
import numpy as np

def main():
    # 示例数据
    documents = [
        "这是一个关于人工智能的文章",
        "机器学习在各个领域都有应用",
        "深度学习技术日新月异",
        "自然语言处理是AI的重要分支",
        "计算机视觉技术发展迅速",
        "这篇文章讲述了经济发展",
        "金融市场波动剧烈",
        "股市投资需要谨慎",
        "经济增长速度放缓",
        "通货膨胀影响市场"
    ]
    
    labels = ['技术'] * 5 + ['经济'] * 5
    
    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(
        documents, labels, test_size=0.3, random_state=42
    )
    
    # 训练分类器
    classifier = EnhancedClassifier(vector_size=100)
    classifier.train(X_train, y_train)
    
    # 进行预测
    y_pred = [classifier.predict(doc) for doc in X_test]
    
    # 获取预测概率（用于ROC曲线）
    y_prob = classifier.classifier.decision_function([
        classifier.text_processor.get_document_vector(doc) for doc in X_test
    ])
    
    # 计算每个类别的准确率
    unique_labels = list(set(labels))
    class_accuracies = []
    for label in unique_labels:
        mask = np.array(y_test) == label
        correct = np.array(y_pred)[mask] == label
        accuracy = correct.mean() if len(correct) > 0 else 0
        class_accuracies.append(accuracy)
    
    # 创建可视化对象
    visualizer = ClassificationVisualizer()
    
    # 生成各种可视化图表
    plots = [
        visualizer.plot_confusion_matrix(y_test, y_pred, unique_labels),
        visualizer.plot_accuracy_bar(class_accuracies, unique_labels),
        visualizer.plot_roc_curve(
            label_binarize(y_test, classes=unique_labels).ravel(),
            y_prob
        ),
        visualizer.plot_wordcloud(' '.join(documents))
    ]
    
    # 保存图表
    filenames = [
        '/Users/liyinghui/Desktop/nlp_class/Gensim/pic/confusion_matrix.png',
        '/Users/liyinghui/Desktop/nlp_class/Gensim/pic/accuracy_bar.png',
        '/Users/liyinghui/Desktop/nlp_class/Gensim/pic/roc_curve.png',
        '/Users/liyinghui/Desktop/nlp_class/Gensim/pic/wordcloud.png'
    ]
    
    visualizer.save_plots(plots, filenames)
    print("可视化图表已保存到当前目录")

if __name__ == '__main__':
    main()